# Author: Zhicheng

import sys
sys.path.append('D:\yzc\Dropbox\private\proj\cuda-convnet-read-only');
from options import *
from util import *
from data import *
import numpy as np
import numpy.random as nr

class LabeledMemoryBatchDataProvider(LabeledDataProvider):
    def __init__(self,data_dir, batch_range, init_epoch=1, init_batchnum=None, dp_params={}, test=False):
        LabeledDataProvider.__init__(self,data_dir,batch_range,init_epoch,
                                     init_batchnum,dp_params,test)
    def get_next_batch(self):
        epoch, batchnum = self.curr_epoch, self.curr_batchnum
        self.load_batch_data()
        self.advance_batch()
        return epoch, batchnum, self.curr_batchdat_dic
#         return epoch, batchnum, [self.curr_batchdat_dic['data'],self.curr_batchdat_dic['labels']]
            
    def load_batch_data(self):
        epoch, batchnum = self.curr_epoch, self.curr_batchnum
        self.curr_batchdat_dic=unpickle(self.get_data_file_name(batchnum))

        
        self.curr_batchdat_dic['data']=np.require((self.curr_batchdat_dic['data']-self.data_mean),
                                                  dtype=np.single,requirements='C')
        self.curr_batchdat_dic['labels']=np.require(self.curr_batchdat_dic['labels'].reshape((1,self.curr_batchdat_dic['data'].shape[1])), 
                                                   dtype=np.single,requirements='C')
#         print 'data shape'
#         print self.curr_batchdat_dic['data'].shape
#         print 'labels shape'
#         print self.curr_batchdat_dic['labels'].shape

# only current batch data is stored in main memory
class ImagenetDataProvider(LabeledMemoryBatchDataProvider):
    def __init__(self,data_dir,batch_range,init_epoch=1,
                 init_batchnum=None,dp_params={},test=False):
        LabeledDataProvider.__init__(self,data_dir,batch_range,init_epoch,
                 init_batchnum,dp_params,test)
        self.curr_batchdat_dic=[]
        self.data_mean=self.batch_meta['data_mean'] # data_mean shape: (width*height*channel,1)
        self.num_colors=3
        self.img_size=256
        
    def get_next_batch(self):
        epoch, batchnum, curr_batchdat_dic = LabeledMemoryBatchDataProvider.get_next_batch(self)
        return epoch, batchnum, [curr_batchdat_dic['data'], curr_batchdat_dic['labels']]
               
    def get_data_dims(self,idx=0):
        return self.img_size**2*self.num_colors if idx==0 else 1
        
    def get_plottable_data(self,data):
        return np.require((data + self.data_mean).T.reshape(data.shape[1], 3, self.img_size, self.img_size).swapaxes(1,3).swapaxes(1,2) / 255.0, dtype=np.single)
    
class CroppedImagenetDataProvider(LabeledMemoryBatchDataProvider):
    def __init__(self,data_dir,batch_range,init_epoch=1,
                 init_batchnum=None,dp_params={},test=False):
        LabeledDataProvider.__init__(self,data_dir,batch_range,init_epoch,
                                     init_batchnum,dp_params,test)
        self.curr_batchdat_dic=[]
        self.data_mean=self.batch_meta['data_mean'] # data_mean shape: (width*height*channel,1)
        self.num_colors=3
        self.img_size=256
        
        self.border_size = dp_params['crop_border']
        self.inner_size = self.img_size - self.border_size*2
        self.multiview = dp_params['multiview_test'] and test
        self.num_views = 5*2
        self.data_mult = self.num_views if self.multiview else 1
        self.num_colors = 3
           
        # self.cropped_data = [n.zeros((self.get_data_dims(), self.data_dic[0]['data'].shape[1]*self.data_mult), dtype=n.single) for x in xrange(2)] 
        # two times for horizontal reflection?
        
        self.batches_generated = 0
        # self.data_mean = self.batch_meta['data_mean'].reshape((3,self.img_size,self.img_size))[:,self.border_size:self.border_size+self.inner_size,self.border_size:self.border_size+self.inner_size].reshape((self.get_data_dims(), 1))

    def get_next_batch(self):
        epoch, batchnum, curr_batchdat_dic = LabeledMemoryBatchDataProvider.get_next_batch(self)
        # cropped = self.cropped_data[self.batches_generated % 2]
        cropped = np.zeros((self.get_data_dims(), curr_batchdat_dic['data'].shape[1]*self.data_mult), dtype=np.single)
        self.__trim_borders(curr_batchdat_dic['data'], cropped)
        # cropped -= self.data_mean
        self.batches_generated += 1
        return epoch, batchnum, [cropped, curr_batchdat_dic['labels']]       
        
    def get_data_dims(self, idx=0):
        return self.inner_size**2 * 3 if idx == 0 else 1

    # Takes as input an array returned by get_next_batch
    # Returns a (numCases, imgSize, imgSize, 3) array which can be
    # fed to pylab for plotting.
    # This is used by shownet.py to plot test case predictions.
    def get_plottable_data(self, data):
        data_mean_center=self.data_mean[:,self.border_size:self.border_size+self.inner_size,self.border_size:self.border_size+self.inner_size]
        return np.require((data + data_mean_center).T.reshape(data.shape[1], 3, self.inner_size, self.inner_size).swapaxes(1,3).swapaxes(1,2) / 255.0, dtype=np.single)

                   
    def __trim_borders(self, x, target):
        y = x.reshape(3, self.img_size, self.img_size, x.shape[1])

        if self.test: # don't need to loop over cases
            if self.multiview:
                start_positions = [(0,0),  (0, self.border_size*2),
                                   (self.border_size, self.border_size),
                                  (self.border_size*2, 0), (self.border_size*2, self.border_size*2)]
                end_positions = [(sy+self.inner_size, sx+self.inner_size) for (sy,sx) in start_positions]
                for i in xrange(self.num_views/2):
                    pic = y[:,start_positions[i][0]:end_positions[i][0],start_positions[i][1]:end_positions[i][1],:]
                    target[:,i * x.shape[1]:(i+1)* x.shape[1]] = pic.reshape((self.get_data_dims(),x.shape[1]))
                    target[:,(self.num_views/2 + i) * x.shape[1]:(self.num_views/2 +i+1)* x.shape[1]] = pic[:,:,::-1,:].reshape((self.get_data_dims(),x.shape[1]))
            else:
                pic = y[:,self.border_size:self.border_size+self.inner_size,self.border_size:self.border_size+self.inner_size, :] # just take the center for now
                target[:,:] = pic.reshape((self.get_data_dims(), x.shape[1]))
        else:
            for c in xrange(x.shape[1]): # loop over cases
                startY, startX = nr.randint(0,self.border_size*2 + 1), nr.randint(0,self.border_size*2 + 1)
                endY, endX = startY + self.inner_size, startX + self.inner_size
                pic = y[:,startY:endY,startX:endX, c]
                if nr.randint(2) == 0: # also flip the image with 50% probability
                    pic = pic[:,:,::-1]
                target[:,c] = pic.reshape((self.get_data_dims(),))
                
# class LabeledMemoryIncoreBatchProvider(LabeledMemoryDataProvider):
#     def __init__(self, data_dir, batch_range, init_epoch=1, init_batchnum=None, dp_params={}, test=False):
#         LabeledMemoryDataProvider.__init__(self, data_dir, batch_range, init_epoch, init_batchnum, dp_params, test)
#         self.data_mean = self.batch_meta['data_mean']
#         self.num_colors = 3
#         self.img_size = 256
#         
#         # Subtract the mean from the data and make sure that both data and
#         # labels are in single-precision floating point.
#         for d in self.data_dic:
#             # This converts the data matrix to single precision and makes sure that it is C-ordered
#             d['data'] = n.require((d['data'] - self.data_mean), dtype=n.single, requirements='C')
#             d['labels'] = n.require(d['labels'].reshape((1, d['data'].shape[1])), dtype=n.single, requirements='C')
# 
#     def get_next_batch(self):
#         epoch, batchnum, datadic = LabeledMemoryDataProvider.get_next_batch(self)
#         return epoch, batchnum, [datadic['data'], datadic['labels']]
# 
#     # Returns the dimensionality of the two data matrices returned by get_next_batch
#     # idx is the index of the matrix. 
#     def get_data_dims(self, idx=0):
#         return self.img_size**2 * self.num_colors if idx == 0 else 1
#     
#     # Takes as input an array returned by get_next_batch
#     # Returns a (numCases, imgSize, imgSize, 3) array which can be
#     # fed to pylab for plotting.
#     # This is used by shownet.py to plot test case predictions.
#     def get_plottable_data(self, data):
#         return n.require((data + self.data_mean).T.reshape(data.shape[1], 3, self.img_size, self.img_size).swapaxes(1,3).swapaxes(1,2) / 255.0, dtype=n.single)
#     
    
